19  O Loop de Treinamento e Otimização

Implementação do ciclo de treinamento, cálculo da função de perda (Cross-Entropy), Backpropagation e configuração do otimizador AdamW.

20 O Loop de Treinamento e Otimização

O “coração” de qualquer sistema de Deep Learning é o loop de treinamento. É neste ciclo iterativo que os parâmetros do modelo (pesos e vieses) são ajustados matematicamente para minimizar o erro entre as previsões do modelo e os dados reais.

Neste capítulo, dissecamos a arquitetura do ciclo de treinamento, a matemática por trás da função de perda Cross-Entropy, a mecânica da Backpropagation e a justificativa para o uso do otimizador AdamW em arquiteturas baseadas em Transformers.

20.1 1. Visão Geral do Ciclo de Treinamento

O treinamento de uma rede neural é um processo de otimização estocástica. O objetivo é encontrar um conjunto de parâmetros \(\theta\) que minimize uma função de custo \(J(\theta)\).

O fluxo de dados ocorre em duas direções: 1. Forward Pass (Propagação): Os dados fluem da entrada para a saída, gerando previsões (logits). 2. Backward Pass (Retropropagação): O erro flui da saída para a entrada, calculando o gradiente da perda em relação a cada parâmetro.

20.1.1 Diagrama de Fluxo do Treinamento

O diagrama abaixo ilustra a interação entre os dados, o modelo, a função de perda e o otimizador.

flowchart TD
    subgraph Data_Preparation [Preparação de Dados]
        Batch[Batch de Treinamento]
        Targets[Labels/Targets]
    end

    subgraph Forward_Pass [Forward Pass]
        Params[(Parâmetros do Modelo)]
        Model[Arquitetura Transformer]
        Logits[Logits / Previsões]
        
        Params --> Model
        Batch --> Model
        Model --> Logits
    end

    subgraph Optimization_Cycle [Ciclo de Otimização]
        LossFn{Cross-Entropy Loss}
        Gradients[Gradientes]
        Optimizer[Otimizador AdamW]
        
        Logits --> LossFn
        Targets --> LossFn
        LossFn -- "loss.backward()" --> Gradients
        Gradients --> Optimizer
        Optimizer -- "step()" --> Params
    end

    style LossFn fill:#f96,stroke:#333,stroke-width:2px
    style Optimizer fill:#69f,stroke:#333,stroke-width:2px
    style Params fill:#eee,stroke:#333,stroke-width:2px

20.2 2. A Função de Perda: Cross-Entropy

Para modelos de linguagem (LLMs), o objetivo é prever o próximo token em uma sequência. Isso é tratado como um problema de classificação multiclasse, onde o número de classes é igual ao tamanho do vocabulário (\(V\)).

A função de perda padrão é a Cross-Entropy (Entropia Cruzada). Ela mede a divergência entre duas distribuições de probabilidade: 1. \(P\) (Distribuição Real): Uma distribuição one-hot representando o token verdadeiro. 2. \(Q\) (Distribuição Predita): As probabilidades geradas pelo modelo (após a aplicação de Softmax nos logits).

A fórmula matemática para um único exemplo é:

\[ H(P, Q) = -\sum_{x} P(x) \log Q(x) \]

No contexto de PyTorch, a classe nn.CrossEntropyLoss combina duas operações para estabilidade numérica: LogSoftmax e NLLLoss (Negative Log Likelihood Loss).

Por que Cross-Entropy? Ela penaliza fortemente previsões confiantes, porém erradas. Se o modelo atribui uma probabilidade baixa ao token correto, o logaritmo dessa probabilidade resulta em um valor negativo grande, aumentando drasticamente a perda.

20.3 3. O Otimizador: AdamW

Embora o Stochastic Gradient Descent (SGD) seja a base, modelos modernos (especialmente Transformers) utilizam variantes adaptativas. O padrão da indústria atual é o AdamW.

20.3.1 Adam vs. AdamW

  • Adam (Adaptive Moment Estimation): Calcula taxas de aprendizado adaptativas para cada parâmetro. Ele armazena médias móveis exponenciais dos gradientes passados (\(m_t\), primeiro momento) e dos gradientes ao quadrado (\(v_t\), segundo momento).
  • O Problema do Adam: A implementação original do Adam aplicava a regularização L2 (Weight Decay) junto com o cálculo do gradiente adaptativo. Isso acoplava a regularização à magnitude dos gradientes, o que muitas vezes levava a uma generalização subótima.
  • A Solução AdamW (Decoupled Weight Decay): O AdamW desacopla o decaimento de peso da atualização do gradiente. O decaimento de peso é aplicado diretamente aos parâmetros após o passo de otimização principal.

A atualização de parâmetros no AdamW segue a lógica:

  1. Atualiza parâmetros baseados nos gradientes adaptativos (Adam clássico).
  2. Aplica decaimento de peso: \(\theta_{t} = \theta_{t} - \eta \lambda \theta_{t-1}\) (onde \(\lambda\) é o fator de decaimento).

Isso permite usar taxas de decaimento de peso mais agressivas sem prejudicar o aprendizado, resultando em modelos que generalizam melhor.

20.4 4. Implementação Técnica

Abaixo apresentamos uma implementação robusta do loop de treinamento utilizando PyTorch.

20.4.1 Estrutura do Código

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.utils.data import DataLoader

def train_epoch(
    model: nn.Module, 
    dataloader: DataLoader, 
    optimizer: torch.optim.Optimizer, 
    device: torch.device,
    clip_grad: float = 1.0
) -> float:
    """
    Executa uma época de treinamento.
    
    Args:
        model: O modelo Transformer.
        dataloader: Iterador de dados em batches.
        optimizer: Otimizador configurado (AdamW).
        device: CPU ou CUDA.
        clip_grad: Valor para recorte de gradiente (Gradient Clipping).
        
    Returns:
        Média da perda (loss) na época.
    """
    model.train() # Coloca o modelo em modo de treinamento (ativa Dropout, etc.)
    total_loss = 0.0
    
    # Definição da Função de Perda
    # ignore_index é usado para ignorar tokens de padding no cálculo da perda
    criterion = nn.CrossEntropyLoss(ignore_index=0) 

    for batch_idx, (inputs, targets) in enumerate(dataloader):
        # 1. Mover dados para o dispositivo (GPU/TPU)
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        # 2. Zerar os gradientes acumulados
        # Essencial: PyTorch acumula gradientes por padrão.
        optimizer.zero_grad(set_to_none=True) 
        
        # 3. Forward Pass
        # O modelo retorna logits de forma (Batch, Seq_Len, Vocab_Size)
        logits = model(inputs)
        
        # 4. Reshape para cálculo da Loss
        # CrossEntropy espera (N, C) ou (N, C, d1...)
        # Achatamos Batch e Seq_Len para tratar cada token como uma instância
        B, T, C = logits.shape
        logits = logits.view(B * T, C)
        targets = targets.view(B * T)
        
        # 5. Cálculo da Perda
        loss = criterion(logits, targets)
        
        # 6. Backward Pass (Backpropagation)
        # Calcula dLoss/dWeights para todos os parâmetros
        loss.backward()
        
        # 7. Gradient Clipping
        # Crucial em Transformers para evitar explosão de gradientes
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad)
        
        # 8. Passo do Otimizador
        # Atualiza os pesos: w = w - lr * grad
        optimizer.step()
        
        total_loss += loss.item()
        
        # (Opcional) Logging intra-época
        if batch_idx % 100 == 0:
            print(f"Batch {batch_idx} | Loss: {loss.item():.4f}")

    return total_loss / len(dataloader)

# Exemplo de Configuração do Otimizador
# model = MyTransformer(...)
# optimizer = AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)

20.5 5. Detalhes Críticos de Arquitetura

20.5.1 Gradient Accumulation (Acumulação de Gradientes)

Se a memória da GPU não suportar o tamanho de batch desejado (ex: 512), usamos acumulação de gradientes. Executamos o forward e backward várias vezes com micro-batches menores, mas só chamamos optimizer.step() após \(N\) passos.

20.5.2 Gradient Clipping

Transformers são profundos e sensíveis. Gradientes podem crescer exponencialmente (exploding gradients), desestabilizando o treino. O clip_grad_norm_ reescala o vetor de gradientes global se sua norma exceder um limiar (geralmente 1.0), mantendo a direção mas limitando a magnitude.

20.5.3 set_to_none=True

Ao zerar gradientes, usar optimizer.zero_grad(set_to_none=True) é ligeiramente mais eficiente em termos de memória do que definir os tensores como zero, pois evita operações de leitura/escrita de memória desnecessárias.

20.6 Resumo do Capítulo

O sucesso do treinamento de um LLM depende do equilíbrio delicado entre a arquitetura do modelo e o algoritmo de otimização. Utilizamos Cross-Entropy para guiar o modelo probabilisticamente e AdamW para navegar na superfície de erro complexa de forma eficiente, desacoplando a regularização da adaptação do passo. A implementação correta do loop, incluindo salvaguardas como Gradient Clipping, é mandatória para a convergência do modelo.